/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <cassert>

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	RC rc;
	PageId eid;

	// Open the index file
	rc = pf.open(indexname.c_str(), mode);
	eid = pf.endPid();
	// Load the root node.
	if (eid > 0){
		BTNonLeafNode root;
		root.read(0, pf);
		rootPid = 0;
		treeHeight = root.getHeight();
	}
	return rc;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    return pf.close();
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	PageId pid;

	if (rootPid == -1){
		// Create a new leaf node and set as the root.
		BTLeafNode root;
		rootPid = pf.endPid();
		treeHeight = 1;
		root.setHeight(treeHeight);
		root.insert(key, rid);
		root.write(rootPid, pf);
	} else {
		BTNonLeafNode node;
		if (treeHeight == 1){
			BTLeafNode root;

			root.read(rootPid, pf);
			if (root.getKeyCount() < BTLeafNode::maxKeyCount){
				// Just insert the key into the root
				root.insert(key, rid);
				root.write(rootPid, pf);
			} else {
			// Split the old root and create a new root
				BTLeafNode sibling;
				BTNonLeafNode newroot;
				PageId oldrootpid;
				int siblingKey;

				pid = pf.endPid();
				root.insertAndSplit(key, rid, sibling, siblingKey);
				sibling.write(pid, pf);
				// Move the old root so that the root stays at 0.
				oldrootpid = pf.endPid();
				root.setNextNodePtr(pid);
				root.write(oldrootpid, pf);
				// Initialize the new root with the old root's and the sibling's pid
				newroot.initializeRoot(oldrootpid, siblingKey, pid);
				newroot.setHeight(++treeHeight);
				newroot.write(rootPid, pf);
			}
		} else {
			PageId npid;
			int ret, nkey;
			
			// Recursively insert the key.
			ret = insertInternal(treeHeight - 1, rootPid, key, rid, nkey, npid);
			if (ret < 0){
				// Split the root
				BTNonLeafNode root, newroot;
				pid = pf.endPid();
				// Move the old root to the end of index
				root.read(rootPid, pf);
				root.write(pid, pf);
				// Initialize new root
				newroot.initializeRoot(pid, nkey, npid);
				newroot.setHeight(++treeHeight);
				newroot.write(rootPid, pf);
				// increase tree height
			}
		}
	}
	return 0;
}

int BTreeIndex::insertInternal(int height, PageId pid, int key, const RecordId& rid,
                               int &newKey, PageId &newPid)
{
	// Child pid
	int ret, nkey;
	PageId cpid, npid;
	BTNonLeafNode node;
	
	node.read(pid, pf);
	assert(node.locateChildPtr(key, cpid) == 0);
	// We must go deeper.
	if (height > 1)
		ret = insertInternal(height - 1, cpid, key, rid, nkey, npid);
	else
	// Write to the leaf node.
		ret = insertLeaf(cpid, key, rid, nkey, npid);
	if (ret != 0){
		// Node is too full
		if (node.getKeyCount() == BTNonLeafNode::maxKeyCount){
			BTNonLeafNode sibling;

			// Allocate a new sibling node
			newPid = pf.endPid();
			// Split this node
			node.insertAndSplit(nkey, npid, sibling, newKey);
			sibling.write(newPid, pf);
			node.write(pid, pf);
			return -1;
		} else {
			// Just insert into the node
			node.insert(nkey, npid);
			node.write(pid, pf);
			return 0;
		}
	}
	return 0;
}

/* Function that actually writes the key to a leaf node. If the leaf is full and splits, newKey and
 * newPid will contain the new sibling key and node.
 * @return: 0 if no split, -1 if split occurred.
 */

int BTreeIndex::insertLeaf(PageId pid, int key, const RecordId& rid,
                           int& newKey, PageId &newPid)
{
	BTLeafNode leaf;

	leaf.read(pid, pf);
	if (leaf.getKeyCount() == BTLeafNode::maxKeyCount){
		// This leaf is full, so split it.
		BTLeafNode sibling;
		PageId spid = leaf.getNextNodePtr();
		leaf.insertAndSplit(key, rid, sibling, newKey);
		newPid = pf.endPid();
		// Sibling contains this node's next pointer
		sibling.setNextNodePtr(spid);
		// Write the sibling to page file
		sibling.write(newPid, pf);
		leaf.setNextNodePtr(newPid);
		leaf.write(pid, pf);
		return -1;
	} else {
		// Just insert the key into the leaf
		leaf.insert(key, rid);
		leaf.write(pid, pf);
		return 0;
	}
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */

RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	BTNonLeafNode internal_node;
	BTLeafNode leaf;
	PageId pid = rootPid;
	int height, eid;
	RC rc;

	height = treeHeight;
	// Iteratively locate the leaf node
	while (--height > 0){
		internal_node.read(pid, pf);
		if ((rc = internal_node.locateChildPtr(searchKey, pid)) < 0)
			return rc;
	}
	// Once we have the leaf, set up the cursor
	leaf.read(pid, pf);
	if ((rc = leaf.locate(searchKey, eid)) == 0){
		cursor.pid = pid;
		cursor.eid = eid;
	}
	return rc;
}

#define MIN_INT		(1 << 31)
void BTreeIndex::verify(void)
{
	IndexCursor cursor;
	locate(MIN_INT, cursor);
	int last;
	int curr;
	RecordId rid;
	readForward(cursor, last, rid);
	while (readForward(cursor, curr, rid) == 0){
		assert(curr >= last);
		last = curr;
	}
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	BTLeafNode leaf;
	RC rc;

	// We've reached the end of the tree.
	if (cursor.pid == -1)
		return RC_END_OF_TREE;
	leaf.read(cursor.pid, pf);
	// Read the entry
	if ((rc = leaf.readEntry(cursor.eid, key, rid)) < 0)
		return rc;
	// Get the pointer to this node's next sibling.
	if (cursor.eid + 1 == leaf.getKeyCount()){
		cursor.pid = leaf.getNextNodePtr();
		cursor.eid = 0;
	} else {
		cursor.eid++;
	}
	return 0;
}
